"""
训练
"""
import numpy as np
import torch
from torch.optim import Adam
from tqdm import tqdm

from chatbot import config
from chatbot.dnn.sort.dataset import dnnsort_data_loader
from chatbot.dnn.sort.siamese import SiameseNetwork

import torch.nn.functional as F

model = SiameseNetwork().to(config.device)
optimizer = Adam(model.parameters(), lr=0.001)


def train(epoch):
    bar = tqdm(enumerate(dnnsort_data_loader), total=len(dnnsort_data_loader))
    loss_list = []
    for idx, (input1, input2, target) in bar:
        optimizer.zero_grad()
        output = model(input1, input2)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        loss_list.append(loss.item())
        bar.set_description("{}\t{}\t{:.6%}".format(epoch, idx, np.mean(loss_list)))
        torch.save(model.state_dict(), config.sort_model_save_path)
        torch.save(optimizer.state_dict(), config.sort_optimizer_save_path)
